package com.example.chat.service;

import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Map;

/**
 * @author qwert
 * @version 1.0.0
 * @create 2025/8/6 22:52
 **/
@Service
public class VectorSearchService {

    @Autowired
    VectorStore vectorStore;

    // 添加单个文档
    public void addDocument(String id, String content, Map<String, Object> metadata) {
        Document document = new Document(content, metadata);
        document.getMetadata().put("id", id);
        vectorStore.add(List.of(document));
    }

    // 批量添加文档
    public void addDocuments(List<Document> documents) {
        vectorStore.add(documents);
    }

    // 相似度搜索
    public List<Document> similaritySearch(String query) {
        return vectorStore.similaritySearch(query);
    }

    // 带过滤的相似度搜索
    public List<Document> similaritySearch(String query, String filter, int topK) {
        return vectorStore.similaritySearch(SearchRequest.builder()
                .query(query)
                .topK(topK)
                .filterExpression(filter)
                .build());
    }

    // 获取所有文档
    public List<Document> getAllDocuments() {
        return similaritySearch("", "", 1000);
    }

    // 按ID删除文档
    public void deleteDocumentById(String id) {
        List<Document> documents = getAllDocuments();
        documents.stream()
                .filter(doc -> id.equals(doc.getMetadata().get("id")))
                .findFirst()
                .ifPresent(doc -> vectorStore.delete(List.of(doc.getId())));
    }

    // 按元数据删除文档
    public void deleteDocumentsByMetadata(String key, String value) {
        List<Document> documents = getAllDocuments();
        List<String> idsToDelete = documents.stream()
                .filter(doc -> value.equals(doc.getMetadata().get(key)))
                .map(Document::getId)
                .toList();

        vectorStore.delete(idsToDelete);
    }

    // 清空向量存储
    public void clearVectorStore() {
        vectorStore.delete(getAllDocuments().stream()
                .map(Document::getId)
                .toList());
    }
}
